#!/bin/bash

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO] $1${NC}"
}

log_warn() {
    echo -e "${YELLOW}[WARN] $1${NC}"
}

log_error() {
    echo -e "${RED}[ERROR] $1${NC}"
}

# 检查命令是否存在
check_command() {
    if ! command -v $1 &> /dev/null; then
        log_error "$1 命令未找到，请先安装"
        exit 1
    fi
}

# 检查 Kudu 服务状态
check_kudu_status() {
    if ! kudu cluster status &> /dev/null; then
        log_error "Kudu 服务未运行"
        exit 1
    fi
}

# 节点管理函数
add_tserver() {
    local node_name=$1
    local node_ip=$2
    local master_ip=$3

    log_info "开始添加 TabletServer: $node_name ($node_ip)"

    # 检查节点连通性
    if ! ping -c 1 $node_ip &> /dev/null; then
        log_error "无法连接到节点 $node_ip"
        return 1
    fi

    # 在远程节点上安装 Kudu
    ssh $node_ip "sudo apt update && sudo apt install -y kudu"

    # 配置远程节点
    ssh $node_ip "sudo tee /etc/kudu/conf/tserver.gflagfile << EOF
--tserver_master_addrs=$master_ip:7051
--fs_wal_dir=/var/lib/kudu/tserver
--fs_data_dirs=/var/lib/kudu/tserver
--rpc_bind_addresses=$node_ip:7050
--webserver_interface=$node_ip
--webserver_port=8050
EOF"

    # 启动远程节点服务
    ssh $node_ip "sudo systemctl daemon-reload && sudo systemctl enable kudu-tserver && sudo systemctl start kudu-tserver"

    # 验证节点加入
    sleep 10
    if kudu cluster status | grep -q $node_ip; then
        log_info "TabletServer $node_name 成功加入集群"
    else
        log_error "TabletServer $node_name 加入集群失败"
        return 1
    fi
}

remove_tserver() {
    local node_name=$1
    local node_ip=$2

    log_info "开始移除 TabletServer: $node_name ($node_ip)"

    # 检查节点是否存在
    if ! kudu cluster status | grep -q $node_ip; then
        log_error "TabletServer $node_name 不在集群中"
        return 1
    fi

    # 停止远程节点服务
    ssh $node_ip "sudo systemctl stop kudu-tserver && sudo systemctl disable kudu-tserver"

    # 清理远程节点数据
    ssh $node_ip "sudo rm -rf /var/lib/kudu/tserver/*"

    log_info "TabletServer $node_name 已成功移除"
}

# 磁盘管理函数
add_disk() {
    local node_name=$1
    local disk_path=$2

    log_info "开始为节点 $node_name 添加磁盘: $disk_path"

    # 检查节点是否存在
    if ! kudu cluster status | grep -q $node_name; then
        log_error "节点 $node_name 不在集群中"
        return 1
    fi

    # 在节点上创建数据目录
    ssh $node_name "sudo mkdir -p $disk_path && sudo chown -R kudu:kudu $disk_path"

    # 更新 Kudu 配置
    ssh $node_name "sudo tee -a /etc/kudu/conf/tserver.gflagfile << EOF
--fs_data_dirs=$disk_path
EOF"

    # 重启节点服务
    ssh $node_name "sudo systemctl restart kudu-tserver"

    # 验证磁盘添加
    sleep 10
    if kudu cluster status | grep -q $disk_path; then
        log_info "磁盘 $disk_path 已成功添加到节点 $node_name"
    else
        log_error "磁盘添加失败"
        return 1
    fi
}

handle_failed_disk() {
    local node_name=$1
    local disk_path=$2

    log_info "开始处理节点 $node_name 的故障磁盘: $disk_path"

    # 检查节点是否存在
    if ! kudu cluster status | grep -q $node_name; then
        log_error "节点 $node_name 不在集群中"
        return 1
    fi

    # 从配置中移除故障磁盘
    ssh $node_name "sudo sed -i '/$disk_path/d' /etc/kudu/conf/tserver.gflagfile"

    # 重启节点服务
    ssh $node_name "sudo systemctl restart kudu-tserver"

    # 卸载故障磁盘
    ssh $node_name "sudo umount $disk_path"

    # 标记磁盘为故障
    ssh $node_name "sudo echo 'DISK_FAILED' > $disk_path/disk_status"

    log_info "故障磁盘 $disk_path 已成功处理"
}

# 主函数
main() {
    # 检查必要命令
    check_command kudu
    check_command ssh
    check_command systemctl

    # 检查 Kudu 状态
    check_kudu_status

    case "$1" in
        "add-tserver")
            if [ $# -ne 4 ]; then
                log_error "用法: $0 add-tserver <node_name> <node_ip> <master_ip>"
                exit 1
            fi
            add_tserver "$2" "$3" "$4"
            ;;
        "remove-tserver")
            if [ $# -ne 3 ]; then
                log_error "用法: $0 remove-tserver <node_name> <node_ip>"
                exit 1
            fi
            remove_tserver "$2" "$3"
            ;;
        "add-disk")
            if [ $# -ne 3 ]; then
                log_error "用法: $0 add-disk <node_name> <disk_path>"
                exit 1
            fi
            add_disk "$2" "$3"
            ;;
        "handle-failed-disk")
            if [ $# -ne 3 ]; then
                log_error "用法: $0 handle-failed-disk <node_name> <disk_path>"
                exit 1
            fi
            handle_failed_disk "$2" "$3"
            ;;
        *)
            log_error "用法: $0 {add-tserver|remove-tserver|add-disk|handle-failed-disk}"
            exit 1
            ;;
    esac
}

# 执行主函数
main "$@" 